package com.china.produce.controller;

import org.springframework.beans.BeansException;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.context.ApplicationContext;
import org.springframework.context.ApplicationContextAware;
import org.springframework.data.redis.core.ReactiveStringRedisTemplate;
import org.springframework.data.redis.core.script.RedisScript;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestMethod;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
import java.time.Instant;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.concurrent.TimeUnit;


/**
 * 令牌限流（桶的最大容量burstCapacity，填充的每分钟流速replenishRate、lua脚本、）
 *
 * @ClassName
 * @Author niujie
 * @Version
 * @Description
 * @CreateTime 2024/7/21
 */
@RequestMapping("/rate")
@RestController
public class RateLimiterController implements ApplicationContextAware {

    @Autowired
    @Qualifier("redisRequestRateLimiterScript")
    private RedisScript<List<Long>> redisRequestRateLimiterScript;

    private ReactiveStringRedisTemplate reactiveStringRedisTemplate;

    private static final String BURST_CAPACITY = "10";

    private static final String REPLENISH_RATE = "3";

    /**
     * 令牌桶限流
     *
     * 桶的最大容量（burstCapacity）、令牌桶填充的流速（replenishRate）
     */
    @RequestMapping(value = "/limited", method = RequestMethod.POST)
    public void limited() throws InterruptedException {

        for (int i = 0; i < 100; i++) {
            TimeUnit.SECONDS.sleep(1);
            Mono<HashMap<String, Object>> allowed = isAllowed();

            allowed.subscribe(data -> {
                data.forEach((key, value) -> {
                    System.out.println("key : " + key + " value: " + value);
                });
            });
        }
    }

    /**
     *
     * @return
     */
    public Mono<HashMap<String, Object>> isAllowed() {
        String prefix = "request_rate_limiter.{test";
        String tokenKey = prefix + "}.tokens";
        String timestampKey = prefix + "}.timestamp";
        List<String> scriptArgs = Arrays.asList(REPLENISH_RATE, BURST_CAPACITY, (Instant.now().getEpochSecond() / 60) + "", "1");

        try {
            Flux<List<Long>> execute = reactiveStringRedisTemplate.execute(redisRequestRateLimiterScript, Arrays.asList(tokenKey, timestampKey), scriptArgs);
            return execute.onErrorResume(throwable -> {
                return Flux.just(Arrays.asList(1L, -1L));
            }).reduce(new ArrayList<Long>(), (longs, l) -> {
                longs.addAll(l);
                return longs;
            }).map(results -> {
                boolean allowed = results.get(0) == 1L;
                Long tokensLeft = results.get(1);
                System.out.println("allowed " + allowed + " tokensLeft " + tokensLeft);
                return getHeaders(-1L);
            });
        }catch (Exception exception) {

            System.out.println(exception.getMessage());
        }

        return Mono.just(getHeaders(-1L));
    }

    public HashMap<String, Object> getHeaders(Long allowed) {

        HashMap<String, Object> map = new HashMap<>();

        map.put("test1", "value1");

        map.put("test2", "value2");

        return map;
    }

    @Override
    public void setApplicationContext(ApplicationContext applicationContext) throws BeansException {

        this.reactiveStringRedisTemplate = applicationContext.getBean(ReactiveStringRedisTemplate.class);
    }
}
